#cheng
import torch 
from  collections import defaultdict
import numpy as np
import random
from typing import TypeVar, Optional, Iterator
import torch.distributed as dist
import math

class ClassAwareSampler(torch.utils.data.sampler.Sampler):
    def __init__(self, dataset, num_replicas: Optional[int] = None, 
                rank: Optional[int] = None, shuffle: bool = True, 
                seed: int = 0, drop_last: bool = False):
        
        if num_replicas is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            num_replicas = dist.get_world_size()
        if rank is None:
            if not dist.is_available():
                raise RuntimeError("Requires distributed package to be available")
            rank = dist.get_rank()
        if rank >= num_replicas or rank < 0:
            raise ValueError(
                "Invalid rank {}, rank should be in the interval"
                " [0, {}]".format(rank, num_replicas - 1))

        self.num_replicas = num_replicas
        self.rank = rank
        self.epoch = 0
        self.drop_last = drop_last
        self.dataset = dataset

        data_dict = defaultdict(list)
        for i, label in enumerate(self._get_labels(dataset)):
            for j in range(label["cls"].shape[0]):
                data_dict[label["cls"][j][0]].append(i)
        num_duplicated_samples = sum([len(v) for v in data_dict.values()])
        class_distribution = [len(v) / num_duplicated_samples for v in data_dict.values()]
        self.sample_indices = []

        frac = 1 / len(data_dict)
        ratios = [frac / p for p in class_distribution]

        for cls_inds, indices in zip(list(data_dict.values()), ratios):
            self.sample_indices += np.random.choice(
                cls_inds, int(len(cls_inds) * frac), replace=True).tolist()

        if self.drop_last and len(self.dataset) % self.num_replicas != 0:  # type: ignore[arg-type]
            # Split to nearest available length that is evenly divisible.
            # This is to ensure each rank receives the same amount of data when
            # using this Sampler.
            self.num_samples = math.ceil(
                (len(self.sample_indices) - self.num_replicas) / self.num_replicas  # type: ignore[arg-type]
            )
        else:
            self.num_samples = math.ceil(len(self.sample_indices) / self.num_replicas)  # type: ignore[arg-type]
        self.total_size = self.num_samples * self.num_replicas
        self.shuffle = shuffle
        self.seed = seed

    def _get_labels(self, dataset):
        return dataset.get_labels()

    def __iter__(self):
        if self.shuffle:
            random.seed(self.seed + self.epoch)
            random.shuffle(self.sample_indices)
        
        if not self.drop_last:
            padding_size = self.total_size - len(self.sample_indices)
            if padding_size <= len(self.sample_indices):
                self.sample_indices += self.sample_indices[:padding_size]
            else:
                self.sample_indices += (self.sample_indices * math.ceil(padding_size / len(self.sample_indices)))[:padding_size]
        else:
            # remove tail of data to make it evenly divisible.
            self.sample_indices = self.sample_indices[:self.total_size]
        assert len(self.sample_indices) == self.total_size

         # subsample
        self.sample_indices = self.sample_indices[self.rank:self.total_size:self.num_replicas]
        assert len(self.sample_indices) == self.num_samples

        return iter(self.sample_indices)
        
    def __len__(self):
        return len(self.sample_indices)
    
    def set_epoch(self, epoch: int) -> None:
        self.epoch = epoch